from typing import Union

import numpy as np

from ril_env.envs.grid_world_v1 import GridWorldEnv

DONT_GO_TYPE = [GridWorldEnv.TYPE_TRAP, GridWorldEnv.TYPE_WALL]


def get_grid(x, y, state: np.ndarray):
    if state.shape[0] > x >= 0 and state.shape[1] > y >= 0:
        return state[x, y]
    else:
        return GridWorldEnv.TYPE_WALL


def grid_world_rule(state: Union[np.ndarray, None], act: int):
    if state is None:
        return act, None, True
    if len(state.shape) == 3:
        state = state[-1]
    mask = [True] * GridWorldEnv.action_space.n
    player = np.argwhere(state == GridWorldEnv.TYPE_PLAYER)[0]

    def get_pos(dx, dy):
        return (player + np.array([dx, dy])).tolist()
    mask[GridWorldEnv.MOVE_UP] = get_grid(*get_pos(-1, 0), state) not in DONT_GO_TYPE
    mask[GridWorldEnv.MOVE_DOWN] = get_grid(*get_pos(1, 0), state) not in DONT_GO_TYPE
    mask[GridWorldEnv.MOVE_LEFT] = get_grid(*get_pos(0, -1), state) not in DONT_GO_TYPE
    mask[GridWorldEnv.MOVE_RIGHT] = get_grid(*get_pos(0, 1), state) not in DONT_GO_TYPE

    return act, mask, True
